from ultralytics import YOLO
import os

def main():
    # 检查数据路径
    assert os.path.exists("datasets/images"), "图片目录不存在"
    assert os.path.exists("datasets/labels"), "标签目录不存在"

    # 加载预训练模型
    model = YOLO("yolo11n-pose.pt")

    model.train(
        data="data.yaml",
        epochs=100,
        batch=16,
        imgsz=640,
        device="cpu",
        workers=0,
        project="TrainData",
        name="exp",  # 使用新名称区分（原exp）
        exist_ok=True,
    )

if __name__ == "__main__":
    main()